-- 兼容sysbench 1.0.20的基准测试脚本
-- 修复查询结果访问方式，支持500万数据且不清理
local table_name = "zbench_cameras"
local brands = {"Canon", "Nikon", "Sony", "Fujifilm", "Panasonic", "Leica", "Olympus", "Pentax"}
local types = {"DSLR", "Mirrorless", "Compact", "Bridge", "Action", "Medium Format"}
local con  -- 线程级连接对象
local default_table_size = 5000000

-- 生成随机字符串
local function random_string(length)
    local chars = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
    local result = ""
    for i = 1, length do
        local idx = math.random(1, #chars)
        result = result .. string.sub(chars, idx, idx)
    end
    return result
end

-- 兼容早期版本的查询结果处理函数
local function get_query_result_value(res)
    -- 尝试多种方式获取查询结果中的值，兼容不同sysbench版本
    if type(res) == "table" then
        -- 尝试数组访问方式
        if res[1] and res[1][1] then
            return res[1][1]
        end
        -- 尝试迭代器方式（最兼容早期版本）
        for row in res:rows() do
            for col in row:cols() do
                return col
            end
        end
    end
    return 0
end

-- 准备阶段
function prepare()
    local drv = sysbench.sql.driver()
    con = drv:connect()

    -- 检查表是否存在（使用兼容的结果处理方式）
    local check_sql = string.format(
        "SELECT COUNT(*) FROM information_schema.tables WHERE table_name = '%s'",
        table_name
    )
    local res = con:query(check_sql)
    local table_exists = get_query_result_value(res) > 0

    if not table_exists then
        -- 创建表
        local create_sql = [[
        CREATE TABLE ]] .. table_name .. [[ (
          id int NOT NULL AUTO_INCREMENT,
          `model_name` varchar(255) NOT NULL,
          `brand` varchar(100) NOT NULL,
          `type` varchar(100) NOT NULL,
          `sensor_type` varchar(100) NOT NULL,
          `sensor_size` varchar(100) NOT NULL,
          `resolution` int NOT NULL,
          `max_aperture` varchar(50) NOT NULL,
          `optical_zoom` int DEFAULT NULL,
          `digital_zoom` int DEFAULT NULL,
          `max_shutter_speed` varchar(50) DEFAULT NULL,
          `iso_range` varchar(100) DEFAULT NULL,
          `video_resolution` varchar(100) DEFAULT NULL,
          `video_fps` varchar(50) DEFAULT NULL,
          `display_size` decimal(3,1) DEFAULT NULL,
          `storage_type` varchar(100) NOT NULL,
          `battery_type` varchar(100) NOT NULL,
          `weight` int NOT NULL,
          `dimensions` varchar(50) NOT NULL,
          `color` varchar(50) NOT NULL,
          `price` decimal(10,2) NOT NULL,
          `release_date` date NOT NULL,
          `warranty_period` int NOT NULL,
          `is_available` tinyint(1) NOT NULL,
          `included_accessories` varchar(255) DEFAULT NULL,
          `sales_volume` int NOT NULL,
          `rating` decimal(2,1) NOT NULL,
          `review_count` int NOT NULL,
          `updated_at` datetime NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
          PRIMARY KEY (id),
          KEY `idx_brand` (`brand`),
          KEY `idx_updated_at_id` (`updated_at`,`id`)
        ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
        ]]
        con:query(create_sql)
        print("表 " .. table_name .. " 创建成功")
    else
        print("表 " .. table_name .. " 已存在")
    end

    -- 计算需要插入的数据量
    local target_size = sysbench.opt.table_size or default_table_size
    local current_size = 0

    if table_exists then
        local count_sql = "SELECT COUNT(*) FROM " .. table_name
        local res = con:query(count_sql)
        current_size = get_query_result_value(res)
    end

    local need_insert = target_size - current_size
    if need_insert <= 0 then
        print("数据量已满足需求（" .. current_size .. "条）")
        con:disconnect()
        return
    end

    -- 批量插入设置
    local batch_size = 1000
    local total_batches = math.ceil(need_insert / batch_size)
    print(string.format("需要插入 %d 条数据，分 %d 批次", need_insert, total_batches))

    if not table_exists then
        con:query("ALTER TABLE " .. table_name .. " DISABLE KEYS")
    end

    local start_time = os.time()
    local inserted = 0

    for batch = 1, total_batches do
        local values = {}
        local current_batch_size = math.min(batch_size, need_insert - inserted)

        for i = 1, current_batch_size do
            local brand = brands[math.random(#brands)]
            local model = brand .. "_" .. random_string(8)
            local type = types[math.random(#types)]

            table.insert(values, string.format(
                "('%s', '%s', '%s', 'APS-C', '23.5x15.6mm', %d, " ..
                "'f/2.8', %d, %d, '1/4000', " ..
                "'100-6400', '4K', '30', 3.2, 'SD', " ..
                "'LP-E6', %d, '120x80x70mm', 'Black', %.2f, '2022-01-15', " ..
                "24, 1, 'Charger, Cable', %d, " ..
                "4.5, %d)",
                model, brand, type,
                math.random(18, 36),
                math.random(3, 10), math.random(2, 5),
                math.random(400, 800),
                math.random(3000, 15000) / 10,
                math.random(500, 5000),
                math.random(100, 2000)
            ))
        end

        local insert_sql = string.format(
            "INSERT INTO %s (model_name, brand, type, sensor_type, sensor_size, resolution, " ..
            "max_aperture, optical_zoom, digital_zoom, max_shutter_speed, iso_range, " ..
            "video_resolution, video_fps, display_size, storage_type, battery_type, weight, " ..
            "dimensions, color, price, release_date, warranty_period, is_available, " ..
            "included_accessories, sales_volume, rating, review_count) " ..
            "VALUES %s",
            table_name,
            table.concat(values, ",")
        )

        con:query(insert_sql)
        inserted = inserted + current_batch_size

        if batch % 10 == 0 then
            local elapsed = os.time() - start_time
            local speed = inserted / elapsed
            local remaining = (total_batches - batch) * batch_size
            local eta = remaining / speed
            print(string.format(
                "已插入 %d/%d 条 (%.2f%%)，速度: %.0f条/秒，预计剩余: %.0f秒",
                inserted, need_insert, (inserted/need_insert)*100, speed, eta
            ))
        end
    end

    if not table_exists then
        con:query("ALTER TABLE " .. table_name .. " ENABLE KEYS")
    end

    local total_time = os.time() - start_time
    print(string.format(
        "数据插入完成，共插入 %d 条，耗时 %d 秒，平均速度: %.0f条/秒",
        inserted, total_time, inserted / total_time
    ))

    con:disconnect()
end

-- 清理阶段（保留数据）
function cleanup()
    print("清理功能已禁用，数据将被保留")
end

-- 生成随机相机数据的辅助函数
local function generate_random_camera_data()
    local brand = brands[math.random(#brands)]
    local model = brand .. "_" .. random_string(8)
    local type = types[math.random(#types)]

    return string.format(
        "('%s', '%s', '%s', 'APS-C', '23.5x15.6mm', %d, " ..
        "'f/2.8', %d, %d, '1/4000', " ..
        "'100-6400', '4K', '30', 3.2, 'SD', " ..
        "'LP-E6', %d, '120x80x70mm', 'Black', %.2f, '2022-01-15', " ..
        "24, 1, 'Charger, Cable', %d, " ..
        "4.5, %d)",
        model, brand, type,
        math.random(18, 36),
        math.random(3, 10), math.random(2, 5),
        math.random(400, 800),
        math.random(3000, 15000) / 10,
        math.random(500, 5000),
        math.random(100, 2000)
    )
end

-- 测试阶段
function event()
    -- 50%概率执行INSERT，50%概率执行UPDATE
    if math.random(1, 2) == 1 then
        -- INSERT操作
        local insert_sql = string.format(
            "INSERT INTO %s (model_name, brand, type, sensor_type, sensor_size, resolution, " ..
            "max_aperture, optical_zoom, digital_zoom, max_shutter_speed, iso_range, " ..
            "video_resolution, video_fps, display_size, storage_type, battery_type, weight, " ..
            "dimensions, color, price, release_date, warranty_period, is_available, " ..
            "included_accessories, sales_volume, rating, review_count) " ..
            "VALUES %s",
            table_name,
            generate_random_camera_data()
        )
        con:query(insert_sql)
    else
        -- UPDATE操作
        local brand = brands[math.random(#brands)]
        local query_sql = string.format(
            "SELECT id, model_name, price FROM %s WHERE brand = '%s' LIMIT 5",
            table_name, brand
        )
        con:query(query_sql)

        local id = math.random(1, sysbench.opt.table_size or default_table_size)
        local update_sql = string.format(
            "UPDATE %s SET price = price * 1.01, sales_volume = sales_volume + 1 WHERE id = %d",
            table_name, id
        )
        con:query(update_sql)
    end
end

-- 线程初始化
function thread_init()
    local drv = sysbench.sql.driver()
    con = drv:connect()
end

-- 线程清理
function thread_done()
    if con then
        con:disconnect()
    end
end

-- 注册回调
sysbench.hooks = {
    prepare = prepare,
    cleanup = cleanup,
    event = event,
    thread_init = thread_init,
    thread_done = thread_done
}
